//	Prior for reconstructing the Equation of State of Dark Energy EoS
//	ref:
//
//  @Jan-19-2017: This file is mainly copied from another projetc of mine, xcos.
//  Some improvemnet were made to make the code more readable.

#include <armadillo>
#include <imcmc/imcmc.hpp>
#include <imcmc/parser++.hpp>
#include "Prior.hpp"

using namespace imcmc;
using namespace imcmc::parser;
using namespace std;

DDE_CPZ::DDE_CPZ(){
    floating = false;
    z = NULL;
    a = NULL;
}

DDE_CPZ::~DDE_CPZ() {

    if( dde_spacing == 0 )
        delete[] a;
    else if( dde_spacing == 1 )
        delete[] z;

}

void DDE_CPZ::Init( std::string& prior_settings ) {

    data_info.GetInfo(prior_settings);

    floating		= Read::Read_Bool_from_File(prior_settings, "dde_wfid_floating");
    nw 				= Read::Read_Int_from_File(prior_settings, "eos_bin_num");

    dde_z_max 		= Read::Read_Double_from_File(prior_settings, "tabulated_w_z_max");
    dde_a_min		= 1./(1.+dde_z_max);

    sigmafid 		= Read::Read_Double_from_File(prior_settings, "dde_cpz_prior_sigmafid");

//	============================================================
//	@ Nov-8, 2016, update 'sigma_mean_w' option names ...
//	sigma_mean_w is the same for either evenly spaced in z or a
//	============================================================
    sigma_mean_w	= Read::Read_Double_from_File(prior_settings, "dde_cpz_prior_sigma_mean_w");

    C	= arma::zeros(nw, nw);
    iC	= arma::zeros(nw, nw);

    dde_spacing = Read::Read_Int_from_File(prior_settings,"dde_spacing");
    if( dde_spacing == 0 )
        a = new double[nw];
    else if( dde_spacing == 1 )
        z = new double[nw];

//	====================================================================
//	NOTE: zbin[] and abin[] defined below is DIFFERENT from that used in
//	eos approximation !!!
//	====================================================================
    double zbin[nw];
    double abin[nw];

    double Delta;
    double xi0;
    double xij, x_plus, x_minus, xbar;

    if( dde_spacing == 0 ) {		//	evenly spaced in a

        ac	= Read::Read_Double_from_File(prior_settings, "dde_cpz_prior_ac");

		if( Read::Has_Key_in_File(prior_settings,"dde_cpz_prior_as") ){
			as	= Read::Read_Double_from_File(prior_settings, "dde_cpz_prior_as");
		}
		else{
			as	= ac;
		}

        double da = (1.0-dde_a_min)/nw;
        for(int i=0; i<nw; ++i) {
            abin[i] = 1.0 - (i+0.5)*da;

            if( floating )
                a[i] = abin[i];
        }

        Delta 	= da;
        xi0		= sigma_mean_w*sigma_mean_w*(1-dde_a_min)/_PI_/ac;
    }
    else if( dde_spacing == 1 ) {	//	evenly spaced in z

        zc 	= Read::Read_Double_from_File(prior_settings, "dde_cpz_prior_zc");

		if( Read::Has_Key_in_File(prior_settings,"dde_cpz_prior_zs") ) {
			zs	= Read::Read_Double_from_File(prior_settings, "dde_cpz_prior_zs");
		}
		else{
			zs	= zc;
		}

        double dz = dde_z_max/nw;
        for(int i=0; i<nw; ++i) {
            zbin[i] = (i+0.5)*dz;

            if( floating )
                z[i] = zbin[i];
        }

        Delta 	= dz;
        xi0		= sigma_mean_w*sigma_mean_w*dde_z_max/_PI_/zc;
    }
    else {
        string err = "\n*** DDE_CPZ::Init() ==> unsupported z-spacing !!!";
        throw runtime_error(err);
    }

    for(int i=0; i<nw; ++i) {
        for(int j=0; j<nw; ++j) {

            if( dde_spacing == 0 ) {

                xij 	= fabs(abin[i] - abin[j]);
                x_plus 	= (xij + Delta)/ac;
                x_minus	= (xij - Delta)/ac;
                xbar	= xij/ac;

                C(i,j) 	= xi0*ac*ac/Delta/Delta*(x_plus*atan(x_plus) + x_minus*atan(x_minus)
                                                 - 2.*xbar*atan(xbar)
                                                 + log(1.+xbar*xbar)
                                                 - 0.5*log(1.+x_plus*x_plus)
                                                 - 0.5*log(1.+x_minus*x_minus) );
            }
            else if( dde_spacing == 1 ) {

                xij 	= fabs(zbin[i] - zbin[j]);
                x_plus 	= (xij + Delta)/zc;
                x_minus	= (xij - Delta)/zc;
                xbar	= xij/zc;

                C(i,j) 	= xi0*zc*zc/Delta/Delta
                          * ( x_plus*atan(x_plus) + x_minus*atan(x_minus)
                              - 2.*xbar*atan(xbar)
                              + log(1.+xbar*xbar)
                              - 0.5*log(1.+x_plus*x_plus)
                              - 0.5*log(1.+x_minus*x_minus) );
            }
            else {
                throw runtime_error("*** DDE_CPZ::Init() ==> error in initialize x_plus or x_minus or xbar");
            }

            if( !floating )
                C(i,j)  = C(i,j) + sigmafid*sigmafid;	//  add sigma_fid^2
        }
    }

//    C.save("debug_cpz_no_avg.txt",arma::arma_ascii);

    iC	= C.i();	//	evaluate the inverse of the covariance matrix

//  make a copy of the original covariance matrix without local-average
    C_original  = C;
    iC_original = iC;

//    iC.save("debug_cpz_no_avg_iC.txt",arma::arma_ascii);

//	add the smoothing matrix into (inv-)covariance matrix
    S	= arma::zeros(nw, nw);
    I	= arma::eye(nw,nw);
    if( floating ) {	//	this part can be further improved, by constructing a matrix A so that w^{fid} = A*w

        for(int i=0; i<nw; ++i) {
            int count=0;

            if( dde_spacing == 0 ) {
                for( int j=0; j<nw; ++j ) {
                    if( fabs(a[j] - a[i]) <= as ) {
                        S(i,j) = 1.0;
                        ++count;
                    }
                }
            }
            else if( dde_spacing == 1 ) {
                for( int j=0; j<nw; ++j ) {
                    if( fabs(z[j] - z[i]) <= zs ) {
                        S(i,j) = 1.0;
                        ++count;
                    }
                }
            }

            for( int j=0; j<nw; ++j )
                S(i,j) = S(i,j)/count;
        }

        IS = I-S;
        iC = IS.t() * iC * IS;

//        iC.save("debug_cpz.txt",arma::arma_ascii);
//        exit(0);
    }

}

double Prior_DDE_CPZ(  	imcmc_double&   param,
                        double&         lndet,
                        double&         chisq,
                        void*           model,
                        void*           data,
                        istate&         state ) {

    state.this_like_is_ok=true;

    lndet = chisq = 0;

    DDE_CPZ *cpz = static_cast<DDE_CPZ*>(data);

    arma::rowvec row_w = arma::zeros<arma::rowvec>(cpz->nw);
    arma::colvec col_w = arma::zeros<arma::colvec>(cpz->nw);

//  those extra high-z w_i will not be constrained by CPZ prior.

    for(int i=0; i<cpz->nw; ++i) {
        string wname= "DDE_w"+Read::IntToString(i);
        row_w[i] = param[wname];
        col_w[i] = row_w[i];
    }

//	Note: inverse of the covariance matrix has been transformed to
//	(I-S)^T * iC * (I-S)

    chisq = arma::as_scalar( row_w * cpz->iC * col_w );

    return -lndet - 0.5*chisq;
}
